#include <string>
#include "graphchi_basic_includes.hpp"

using namespace graphchi;

#define NITERS 1000

/**
  * Type definitions. Remember to create suitable graph shards using the
  * Sharder-program. 
  */

typedef struct INFO{	
	int L;
	int sigma;
	double delta;
	double Cb;
} Info;

class Label{
public:
	int id1;
	int id2;
	Label()
	{
		id1 = -1;
		id2 = -1;
	}
	int Neighbor_ID(int id)
	{
		if(id1 == id)
			return id2;
		if(id2 == id)
			return id1;
		return -1;
	}
	void Add_ID(int id)
	{
		if(id1 == -1)
		{
			id1 = id;
			return;
		}
		if(id2 == -1)
		{
			id2 = id;
			return;
		}
	}
};

typedef int VertexDataType;
typedef Label EdgeDataType;

/**
  * GraphChi programs need to subclass GraphChiProgram<vertex-type, edge-type> 
  * class. The main logic is usually in the update function.
  */
struct BetweennessProgram : public GraphChiProgram<VertexDataType, EdgeDataType> {
    
    int INFINITY;
	int N;
	int L;
	int S;
	int M;
	int Step;
	
	Info *info;
	BetweennessProgram()
	{
		L = 0;
		S = 0;
		M = 0;
		Step = 0;
	}
	
    /**
     *  Vertex update function.
     */
    void update(graphchi_vertex<VertexDataType, EdgeDataType> &vertex, graphchi_context &gcontext)
    {
		if (gcontext.iteration == 0) {
			for(int i=0; i < vertex.num_edges(); i++) 
			{
				Label label = vertex.edge(i)->get_data();
				label.Add_ID(vertex.id());
				vertex.edge(i)->set_data(label);
			}
		}
		else
		{			
			int v = vertex.id();
			if(Step == 0)
			{
				if(info[v].L == L)
				{
					for(int i=0; i < vertex.num_edges(); i++)
					{
						int w = vertex.edge(i)->get_data().Neighbor_ID(vertex.id());
						if(info[w].L < 0)
						{
							info[w].L = L+1;
							M++;
						}
					}
					for(int i=0; i < vertex.num_edges(); i++)
					{
						int w = vertex.edge(i)->get_data().Neighbor_ID(vertex.id());
						if(info[w].L == L+1)
							info[w].sigma = info[w].sigma + info[v].sigma;
					}
				}
			}
			if(Step == 1)
			{
				if(info[v].L == L)
				{
					for(int i=0; i < vertex.num_edges(); i++)
					{
						int w = vertex.edge(i)->get_data().Neighbor_ID(vertex.id());
						if(L == info[w].L+1)
							info[w].delta = info[w].delta 
							+ info[w].sigma/(double)(info[v].sigma)*(1+info[v].delta);
					}
					if (v != S)
						info[v].Cb = info[v].Cb + info[v].delta;
				}
			}
		}
    }
    
    /**
     * Called before an iteration starts.
     */
    void before_iteration(int iteration, graphchi_context &gcontext) {
		if (gcontext.iteration == 0) {
			N = gcontext.nvertices;
			INFINITY = N;				
			info = new Info[N];
			for(int i = 0;i < N;i++)
			{
				info[i].sigma = 0;
				info[i].Cb = 0;
				info[i].delta = 0;
			}
		}
		else
			if(Step == 0 && L == 0)
			{
				for(int i = 0;i < N;i++)
				{
					if (i == S)
					{
						info[i].L = 0;
						info[i].sigma = 1;
						M++;
					}
					else
					{
						info[i].L = -1;
						info[i].sigma = 0;
					}
				}
			}
		gcontext.set_last_iteration(NITERS);
    }
    
    /**
     * Called after an iteration has finished.
     */
    void after_iteration(int iteration, graphchi_context &gcontext) 
    {
		if (gcontext.iteration != 0)
		{
			if(Step == 0)
			{
				if(M == N)
				{
					Step = 1;
					for(int i = 0;i < N;i++)
						info[i].delta = 0;
					M = 0;
				}
				L++;
			}
			else
			{
				if(L == 0)
				{
					Step = 0;
					S++;
				}
				else
					L--;
			}
		}
		if(S == N)
			gcontext.set_last_iteration(iteration);
    }
    
    /**
     * Called before an execution interval is started.
     */
    void before_exec_interval(vid_t window_st, vid_t window_en, graphchi_context &gcontext) {        
    }
    
    /**
     * Called after an execution interval has finished.
     */
    void after_exec_interval(vid_t window_st, vid_t window_en, graphchi_context &gcontext) {        
    }
    
};

int main(int argc, const char ** argv) {
    /* GraphChi initialization will read the command line 
       arguments and the configuration file. */
    graphchi_init(argc, argv);    
    /* Metrics object for keeping track of performance counters
       and other information. Currently required. */
    metrics m("betweenness");    
    /* Basic arguments for application */
    std::string filename = get_option_string("file");  // Base filename
    bool scheduler       = get_option_int("scheduler", 0); // Whether to use selective scheduling    
    /* Detect the number of shards or preprocess an input to create them */
    int nshards          = convert_if_notexists<EdgeDataType>(filename, get_option_string("nshards", "auto"));    
    /* Run */
    BetweennessProgram program;
    graphchi_engine<VertexDataType, EdgeDataType> engine(filename, nshards, scheduler, m); 
    engine.run(program, NITERS);
    /* Output top ranked vertices */
    for(int i=0; i < program.N; i++) {
        std::cout << i << "\t" << program.info[i].Cb / 2<< std::endl;
    }
    return 0;
}
